import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  
import os
'''守恒格式测试程序。修改好了边界条件'''

'''输入速度，计算涡度。采用中心差分方法。
函数本义是通用的，但此程序中由于没有更新u,v在中间区域的值，所以就在设初值时用了一次
注意，左侧指标是j，右侧指标是i
最上方：[-1, 1:-1]
'''
def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega


'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream(omega_int, dx, dy, tol=1e-5, psi_init=None):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    change = 1.
    iter = 0
    while change > tol:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        '''
        如果边界上流函数一直为零,初始时设成0就好,后面不动它
        顶盖第二类边条件(收敛性存在问题):psi_new[-1, 1:-1]=psi_new[-2, 1:-1]+dy
        '''
        #psi_new[-1, :]=psi_new[-2, :]   #顶部
        psi_new[-1, :]=0
        psi_new[0, :]=0  #底部
        psi_new[1:-1,0]=psi_new[1:-1,-2]    #周期条件：左侧边界点用右侧内点赋值
        psi_new[1:-1,-1]=psi_new[1:-1,1]    #周期条件：右侧边界点用左侧内点赋值
        
        change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值

        psi_init = psi_new
        iter += 1   
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    return psi_init, end_signal




#边界涡度处理，采用无滑移条件，顶盖速度恒为1，其余三壁速度恒为零.(太耍赖了，用psi和u,v)
def apply_bnd_omega_(psi, u, v, omega):
    #左右边界周期条件
    omega[:,0]=omega[:,-2]
    omega[:,-1]=omega[:,1]

    #顶部底部涡量直接设为零
    omega[0, 1:-1] =0
    omega[-1, 1:-1] =0

#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
    #边界
    u[:,0]=u[:,-2]
    u[:,-1]=u[:,1]
    #u[-1, 1:-1] = u_up
    u[-1, 1:-1]=u[-2, 1:-1]
    u[0, 1:-1] = 0

    v[:,0]=v[:,-2]
    v[:,-1]=v[:,1]
    v[-1, 1:-1] = 0 
    #v[0, 1:-1] = v[1, 1:-1] 
    v[0, 1:-1]=0
    return u,v

#对流算符
def convection_operator(myfun,u,v):
    #原来的非守恒形式：
    #uw_x = u[1:-1,1:-1] / 2 / dx  * (omega[1:-1, 2:] - omega[1:-1, :-2])
    #试着改成守恒形式吧！
    uw_x =  (u[1:-1, 2:]*myfun[1:-1, 2:] -u[1:-1, :-2]*myfun[1:-1, :-2])/ 2 / dx
    vw_y =  (v[2:, 1:-1]*myfun[2:, 1:-1] - v[:-2, 1:-1]*myfun[:-2, 1:-1])/2 / dy

    return uw_x,vw_y

#扩散算符
def diff_operator(myfun,nu):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x,diff_y

#鞘损失算符
def sheath_loss_op(myfun,alpha_sh_y,f_bg):
    myfun_loss=np.zeros_like(myfun)
    for j in range(1,Ny-1):
        for i in range(1,Nx-1):
            myfun_loss[j][i]=-alpha_sh_y[j]*(myfun[j][i]-f_bg)
    return myfun_loss[1:-1, 1:-1]

#漂移驱动
def drive_term_omega(n_density,beta):
    linear_flag=False
    if linear_flag:
        drt=-beta*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
        return drt
    else:
        drt2=np.zeros_like(n_density)
        j_up=-1
        n_log1=np.log(n_density[1:j_up,2:])
        n_log2=np.log(n_density[1:j_up,:-2])
        #drt=beta*(n_log[2:,1:-1]-n_log[:-2,1:-1])/2/dy
        drt2[1:j_up,1:-1]=-beta*(n_log1-n_log2)/2/dy
        return drt2[1:-1,1:-1]


#涡度向前演化
def forward_omega_FTCS(omega,n_density,u,v):
    uw_x,vw_y =convection_operator(omega,u,v)
    diff_x,diff_y=diff_operator(omega,prandtl)  #扩散项用普朗特数
    #psi_sheath_loss=sheath_loss_op(psi,alpha_sh_y,f_bg=0)  #涡度损失项先不要了
    omega_drive_term=drive_term_omega(n_density,Ra*prandtl)
    return omega[1:-1, 1:-1] + dt * (-uw_x  -vw_y+omega_drive_term + diff_y + diff_x)

#密度向前演化
def forward_n_FTCS(n_density,u,v):
    uw_x,vw_y =convection_operator(n_density,u,v)
    diff_x,diff_y=diff_operator(n_density,1)    #归一化的密度扩散系数是1
    n_sheath_loss=sheath_loss_op(n_density,alpha_sh_y,n_0)
    #n_src=n_source_opr(n_density,n_init0_y,nv_sy)  #左侧源项不要了
    return n_density[1:-1, 1:-1] + dt * (-uw_x -vw_y+n_sheath_loss + diff_y + diff_x)

def apply_bnd_n_density(n_density):
    #顶边第二类条件
    n_density[-1,:]=n_density[-2,:]
    #底边第一类条件
    n_density[0,:]=n_down

    #左右边界周期条件
    n_density[:,0]=n_density[:,-2]
    n_density[:,-1]=n_density[:,1]


def diag_n_y_sec(n_density,out_iter):
    
    subfolder='n_density_y_section' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    plt.plot(y,n_density[:,Nx//2])
    plt.xlabel('y')  
    plt.ylabel('n_density') 
    plt.title('Re={},t={}'.format(Re,t_diag))  
    plt.grid(True)
    # 定义文件名，包含idiag的值  
    filename_base = 'n_density_y_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_uv_quiver(u,v,xx,yy,out_iter):
    
    subfolder='uv_quiver' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    # 使用quiver函数绘制矢量图  
    Q = plt.quiver(xx, yy, u, v, units='width', scale=10, color='b')  
    plt.xlabel('X')  
    plt.ylabel('Y')  
    plt.title('2D Vector Field , Re={},t={}'.format(Re,t_diag))  
    # 定义文件名，包含idiag的值  
    filename_base = 'uv_quiver_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_n_pcolor(n_density,out_iter):
    
    subfolder='n_pcolor' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    fig, ax = plt.subplots()
    n_difference=n_density-n_init0_matrix
    # 绘制伪彩色图
    plt.streamplot(xx, yy, u, v, linewidth=2)
    plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 定义文件名，包含idiag的值  
    filename_base = 'n_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_psi_3d(psi,out_iter):
    
    subfolder='psi_3d' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    fig = plt.figure()  
    ax = fig.add_subplot(111, projection='3d')  
    
    # 绘制三维曲面  
    ax.plot_surface(xx, yy, psi, cmap='viridis', linewidth=0)  # cmap参数用于设置颜色映射  
    
    # 添加标题和坐标轴标签  
    ax.set_title('3D Plot of psi,t={:},Ra={:}'.format(t_diag,Ra))  
    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('psi')  
    # 定义文件名，包含idiag的值  
    filename_base = 'psi_3d_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

# 给定的网格步长
dx = 0.02  # 例如，我们希望每个网格点之间的距离是 0.01
dy = 0.02  # 同上

# 给定的模拟区域长度
Lx = 2.0  # 模拟区域在 x 方向的长度
Ly = 2.0  # 模拟区域在 y 方向的长度
Ly_in=Ly/2  # 内部
Ly_out=Ly-Ly_in #外部，相当于delta

# 计算网格点数量
Nx = int(Lx / dx) + 1  # 加 1 是为了确保网格点覆盖整个模拟区域，包括边界
Ny = int(Ly / dy) + 1  # 同上

# 打印结果
print(f"Number of grid points in x-direction (Nx): {Nx}")
print(f"Number of grid points in y-direction (Ny): {Ny}")
u_up=0  #顶盖速度
u_down=0
n_up=n_0=1  #顶盖密度
n_down=1.5    #底边密度
Delta_n=n_down-n_0

x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的。
dx = x[-1] - x[-2]  #网格长
y = np.linspace(0, Ly, Ny)
dy = y[-1] - y[-2]
#nu = 0.001
#D_n=0.001    #密度扩散
alpha_sh0=1
beta=0.1
prandtl=0.1 #Pr,普朗特数
Ra=200000 #Ra*,瑞利数

dt = 0.00002  #时间步长
ntime=10000  #时间步数
ndiag=100 #每隔多少时间输出一次
#Re=1/nu
Re=Ra
xx, yy = np.meshgrid(x, y)  
'''meshgrid生成的矩阵由Ny个长Nx的数列构成，矩阵形状为Ny*Nx。这在索引时造成的顺序是xx[j][i]，如果i是x向指标、j是y向指标的话。
这样，指标数是(i,j)的格点，索引为[j][i]，空间绝对位置(xx[j][i],yy [ j ][i])
'''

if __name__ == '__main__':
    #初始条件设置
    #设置初始流函数场
    psi=np.zeros_like(yy)
    # phi0_y=0.1*sigmoid_lm1_r0(y,c=Ly/2,k=10/Ly)
    # for i in range(0,Nx):
    #     for j in range(0,Ny):
    #         psi[j][i]=phi0_y[j]

    #设置初始速度场
    u = np.zeros_like(xx)
    v = np.zeros_like(yy)
    #u,v=psi_to_vel(psi,u,v,dx,dy)   #用流函数设置
    u[-1, 1:-1] = u_up    #最上方速度。这里设置速度时排除了首尾格点。numpy可以一个数组等于一个数，这样把这个数赋值给这个数组

    # #初始均匀流
    # for i in range(0,Nx):
    #     for j in range(0,Ny):
    #         u[j][i]=yy[j][i]*u_up

    #设置初始密度场
    tmp_alpha=np.sqrt(alpha_sh0)
    tmp_atad=tmp_alpha*np.tanh(tmp_alpha*Ly_out)
    A_2=1/(1+tmp_atad)
    A_1=-tmp_atad*A_2
    B_2=-np.tanh(tmp_alpha*Ly_out)*A_2
    n_init0_y=n_0*y
    for j in range(0,Ny):
        if y[j]<Ly_in:
            n_init0_y[j]=n_0+Delta_n*(1+A_1*y[j])
        else:
            n_init0_y[j]=n_0+Delta_n*(A_2*np.cosh(tmp_alpha*(y[j]-Ly_in))+B_2*np.sinh(tmp_alpha*(y[j]-Ly_in)))
    n_init0_matrix=np.zeros_like(yy)
    for i in range(0,Nx):
            for j in range(0,Ny):
                n_init0_matrix[j][i]=n_init0_y[j]
    n_load_flag=False
    if n_load_flag:
        n_density=np.load('n_3_1.npy')
    else:
        n_density=np.copy(n_init0_matrix)
        # 设定噪声的标准差（std_dev），根据需要调整
        std_dev = 0.001  # 这是一个较小的噪声幅度
        #noise = np.random.normal(scale=std_dev, size=n_density.shape)
        noise=np.random.uniform(-std_dev, std_dev, size=n_density.shape)
        for i in range(0,Nx):
            for j in range(0,Ny*2//3):
                n_density[j][i]=n_density[j][i]*abs(1+noise[j][i])

    #设置alpha_sh_y[j]=alpha_sh(y)
    alpha_sh_y=np.zeros_like(y)
    for j in range(0,Ny):
        #alpha_sh_y[j]=alpha_sh0*H0_right(y[j],Ly/2,Ly/20)
        alpha_sh_y[j]=alpha_sh0*np.heaviside((y[j]-Ly_in),0.5) #阶跃函数,右边系数是中间恰好等于处的值

    omega = np.zeros_like(xx)
    omega[1:-1, 1:-1] = vel_to_omega(u, v, dx, dy)  #计算涡度场。依然排除首尾格点
    psi, _ = compute_stream(omega[1:-1, 1:-1], dx, dy)  #omega只需要中间(Nx-2)*(Ny-2)小方块的信息
    u,v=psi_to_vel(psi,u,v,dx,dy)
    apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。

    #进入主循环
    end_signal = False  #终止判断条件。可以再加一个判断条件。
    out_iter = 0    
    #while end_signal==False and out_iter<1000:
    while  out_iter<ntime:
        if out_iter % ndiag==0:
            diag_n_y_sec(n_density,out_iter)    #截面密度分布
            diag_n_pcolor(n_density,out_iter)
            #diag_psi_3d(psi,out_iter)
            print(out_iter)

        out_iter+=1
        omega[1:-1, 1:-1] = forward_omega_FTCS(omega,n_density,u,v)  #时间向前步进。此时边界涡量会向中心扩散
        n_density[1:-1, 1:-1]=forward_n_FTCS(n_density,u,v)
        psi, end_signal = compute_stream(omega[1:-1, 1:-1], dx, dy, psi_init=psi, tol=1e-5)  #泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。
        #需要新增代码：根据流函数场计算速度场。
        u,v=psi_to_vel(psi,u,v,dx,dy)
        apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。u,v仅用到边界值，其不在程序中改变。因此这里的u,v并不是全流场速度。
        apply_bnd_n_density(n_density)

        
        
    t_total=out_iter*dt #总时间

    #绘制流场
    flag_streamplot=True
    if flag_streamplot:
        n_difference=n_density-n_init0_matrix
        plt.streamplot(xx, yy, u, v, linewidth=2)
        plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), shading='auto')
        plt.colorbar(label='n')
        plt.show()


    flag_pcolor = True
    if flag_pcolor:
        fig, ax = plt.subplots()
        n_difference=n_density-(n_0+Delta_n)*np.ones_like(n_density)
        # 绘制伪彩色图
        cs = ax.pcolor(xx, yy, n_difference, shading='auto', cmap=plt.get_cmap('Spectral_r'))
        # 添加标题
        ax.set_title('Re={}, t={}'.format(Re, t_total), fontsize=16)
        # 添加颜色条
        cbar = fig.colorbar(cs, orientation='vertical')
        # 设置坐标轴的比例相同，使得图形保持矩形
        ax.set_aspect('equal')
        plt.show()
        

    #绘制三维图
    flag_axe3D=True
    if flag_axe3D:
        # 创建三维图形和坐标轴  
        fig = plt.figure()  
        ax = fig.add_subplot(111, projection='3d')  
        
        # 绘制三维曲面  
        ax.plot_surface(xx, yy, psi, cmap='viridis', linewidth=0)  # cmap参数用于设置颜色映射  
        
        # 添加标题和坐标轴标签  
        ax.set_title('3D Plot of ')  
        ax.set_xlabel('X')  
        ax.set_ylabel('Y')  
        ax.set_zlabel('psi')  
        
        # 显示图形  
        plt.show()

    #绘制截面速度图
    flag_u_ysection=True
    if flag_u_ysection:
        #增设画图部分。
        plt.plot(y,u[:,Nx//2])
        plt.title('Re={},t={}'.format(Re,t_total))  
        plt.show()

